{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# 03 Softmax\nSoftmax kernel \u4f18\u5316\u793a\u4f8b"}, {"cell_type": "code", "metadata": {}, "source": "import torch\nimport triton\nimport triton.language as tl\n\n@triton.jit\ndef softmax_kernel(X_ptr, Y_ptr, N: tl.constexpr):\n    pid = tl.program_id(0)\n    if pid < N:\n        x = tl.load(X_ptr + pid)\n        tl.store(Y_ptr + pid, torch.exp(x))  # \u7b80\u5355 exp placeholder\n\nN = 16\nx = torch.randn(N, device='cuda')\ny = torch.empty_like(x)\nsoftmax_kernel[(N,)](x, y, N=N)\nprint(\"y =\", y)", "outputs": [], "execution_count": null}], "metadata": {"title": "Softmax"}, "nbformat": 4, "nbformat_minor": 5}